package io.github.handyang.algorithms.search;

import lombok.Getter;

import java.util.LinkedList;
import java.util.List;

public class BST<K extends Comparable<K>, V> implements SymbolTable<K, V> {
    @Getter
    private TreeNode<K, V> root;

    private int size;

    @Override
    public void put(K key, V value) {
        if (root == null) {
            root = new TreeNode<>(key, value);
            size++;
            return;
        }

        TreeNode<K, V> cursor = this.root;
        while (true) {
            int compare = cursor.getKey().compareTo(key);
            if (compare == 0) {
                cursor.setValue(value);
                return;
            } else if (compare > 0) {
                if (cursor.getLeft() == null) {
                    cursor.setLeft(new TreeNode<>(key, value));
                    size++;
                    return;
                } else cursor = cursor.getLeft();
            } else {
                if (cursor.getRight() == null) {
                    cursor.setRight(new TreeNode<>(key, value));
                    size++;
                    return;
                } else cursor = cursor.getRight();
            }
        }
    }

    @Override
    public V get(K key) {
        TreeNode<K, V> cursor = this.root;

        while (cursor != null) {
            int compare = cursor.getKey().compareTo(key);
            if (compare == 0) {
                return cursor.getValue();
            } else if (compare > 0) {
                cursor = cursor.getLeft();
            } else {
                cursor = cursor.getRight();
            }
        }
        return null;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public Iterable<K> keys() {
        List<K> keys = new LinkedList<>();
        allKey(keys, this.root);

        return keys;
    }

    @Override
    public void clear() {
        this.root = null;
        this.size = 0;
    }

    private void allKey(List<K> keys, TreeNode<K, V> root) {
        if (root == null) return;

        keys.add(root.getKey());

        allKey(keys, root.getLeft());
        allKey(keys, root.getRight());
    }

}
